# asm_feature_extractor.py
import os
import csv
import numpy as np
import joblib
from time import time
from asm_to_kg import process_asm_to_graph_text
from kg_to_bert import process_graph_text_with_bert

# 接口：加载标签
def load_labels(label_file):
    """
    加载标签数据

    Args:
        label_file: 标签文件路径

    Returns:
        labels: 标签字典
    """
    labels = {}
    with open(label_file, 'r') as f:
        reader = csv.DictReader(f)
        for row in reader:
            labels[row['Id']] = int(row['Class'])
    return labels

# 接口：处理ASM文件生成特征向量
def process_asm_file(model_path, asm_path):
    """
    处理单个ASM文件，返回BERT特征向量

    Args:
        model_path: BERT模型路径
        asm_path: 汇编文件路径

    Returns:
        bert_output: BERT特征向量
    """
    # 生成知识图谱文本
    graph_text = process_asm_to_graph_text(asm_path)
    
    # 获取BERT输出向量
    bert_output = process_graph_text_with_bert(model_path, graph_text)
    
    return bert_output.flatten()

# 接口：提取所有样本特征向量生成检查点
def extract_features(model_path, asm_dir, label_file, cache_file, force_recompute=False):
    """
    提取所有样本的特征向量（使用缓存机制）

    Args:
        model_path: BERT模型路径
        asm_dir: 汇编文件夹路径
        label_file: 标签文件路径
        cache_file: 缓存文件路径
        force_recompute: 是否强制重新计算特征（默认为False）

    Returns:
        cache_data: 特征和标签的缓存数据
    """
    if not force_recompute and os.path.exists(cache_file):
        print(f"加载缓存特征: {cache_file}")
        return joblib.load(cache_file)
    
    # 获取标签和文件列表
    labels = load_labels(label_file)
    file_ids = [f.replace('.asm', '') for f in os.listdir(asm_dir) if f.endswith('.asm')]
    
    # 准备数据存储结构
    features = []
    targets = []
    valid_ids = []
    failed_ids = []
    
    print(f"开始处理 {len(file_ids)} 个ASM文件...")
    start_time = time()
    
    for i, file_id in enumerate(file_ids):
        if file_id not in labels:
            continue
            
        asm_path = os.path.join(asm_dir, file_id + '.asm')
        print(f"[{i+1}/{len(file_ids)}] 处理: {file_id}")
        
        try:
            # 处理ASM文件并获取特征
            feature_vector = process_asm_file(model_path, asm_path)
            
            if feature_vector is not None and len(feature_vector) > 0:
                features.append(feature_vector)
                targets.append(labels[file_id])
                valid_ids.append(file_id)
            else:
                print(f"  ⚠️ 无效特征: {file_id}")
                failed_ids.append(file_id)
        except Exception as e:
            print(f"  ❌ 处理失败: {file_id} - {str(e)}")
            failed_ids.append(file_id)
    
    # 转换为numpy数组
    features = np.array(features)
    targets = np.array(targets)
    
    # 保存到缓存
    cache_data = {
        'features': features,
        'targets': targets,
        'valid_ids': valid_ids,
        'failed_ids': failed_ids
    }
    joblib.dump(cache_data, cache_file)
    
    print(f"特征提取完成! 用时: {time()-start_time:.2f}秒")
    print(f"有效样本: {len(valid_ids)}, 失败样本: {len(failed_ids)}")
    
    return cache_data